package org.huangrui.spark.java.sql;

import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.expressions.Aggregator;
import org.huangrui.spark.java.bean.CityCount;
import org.huangrui.spark.java.bean.MyCityRemarkBuffer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
 * @Author hr
 * @Create 2024-10-21 7:19
 */
public class MyCityRemarkUDAF extends Aggregator<String, MyCityRemarkBuffer, String> {
    @Override
    // 缓冲区初始化数据
    public MyCityRemarkBuffer zero() {
        return new MyCityRemarkBuffer(0L, new HashMap<>());
    }

    @Override
    // TODO 将函数的输入值和缓冲区的数据进行聚合处理
    public MyCityRemarkBuffer reduce(MyCityRemarkBuffer b, String a) {
        b.setCount(b.getCount() + 1);
        Map<String, Long> cityMap = b.getCityMap();
        if (cityMap.containsKey(a)) {
            cityMap.put(a, cityMap.get(a) + 1);
        } else {
            cityMap.put(a, 1L);
        }
        b.setCityMap(cityMap);
        return b;
    }

    @Override
    // TODO 合并缓冲区数据
    public MyCityRemarkBuffer merge(MyCityRemarkBuffer b1, MyCityRemarkBuffer b2) {
        b1.setCount(b1.getCount() + b2.getCount());
        Map<String, Long> b1CityMap = b1.getCityMap();
        Map<String, Long> b2CityMap = b2.getCityMap();
        /*
           c1 : { beijing : 10, tianjin : 20, baoding : 30 }
           c2 : { beijing : 40, tianjin : 50, sjz : 60 }
           ------------------------------------------------
           c3 : { beijing : 50, tianjin : 70, baoding : 30, sjz : 60 }

           1. 将 map1 保持不变
           2. 对 map2 进行遍历
           3. 如果 map2 中的 key 在 map1 存在，那么合并数据
           4. 如果 map2 中的 key 不在 map1 存在，那么直接添加即可
         */
        for (String key : b2CityMap.keySet()) {
            final Long v1 = b1CityMap.get(key);
            final Long v2 = b2CityMap.get(key);
            if (v1 == null) {
                b1CityMap.put(key, v2);
            } else {
                b1CityMap.put(key, v1 + v2);
            }
        }
        b1.setCityMap(b1CityMap);
        return b1;
    }

    @Override
    // 将统计的结果生成字符串信息
    public String finish(MyCityRemarkBuffer buffer) {
        StringBuilder sb = new StringBuilder();
        Long total = buffer.getCount();
        Map<String, Long> cityMap = buffer.getCityMap();
        ArrayList<CityCount> cityCounts = new ArrayList<>();
        cityMap.forEach((k, v) -> {
            cityCounts.add(new CityCount(k, v));
        });
        // TODO 排序
        Collections.sort(cityCounts);

        CityCount cityCount = cityCounts.get(0);
        long pc0 = cityCount.getCount() * 100 / total;  // 10 * 100/20 => 50
        sb.append(cityCount.getCityName()).append(" ").append(pc0).append("% ");

        cityCount = cityCounts.get(1);
        long pc1 = cityCount.getCount() * 100 / total;  // 10 * 100/20 => 50
        sb.append(cityCount.getCityName()).append(" ").append(pc1).append("% ");

        if ( cityCounts.size() > 2 ) {
            sb.append("其他 ").append(100 - pc0 - pc1).append("%");
        }
        return sb.toString();
    }

    @Override
    // 缓冲区的编码操作，用于序列化和反序列化数据结构
    public Encoder<MyCityRemarkBuffer> bufferEncoder() {
        return Encoders.bean(MyCityRemarkBuffer.class);
    }

    @Override
    // 输出的编码操作，用于序列化和反序列化函数返回的数据结构
    public Encoder<String> outputEncoder() {
        return Encoders.STRING();
    }
}
